
import csv
import sys
import argparse
from pathlib import Path
from typing import List, Tuple
from PIL import Image
import fitz  # PyMuPDF

# ====== Config ======
OUTPUT_DIR = "xfdata/converted_pngs"           # where PNGs go
OUTPUT_CSV = "xfdata/test_converted_index.csv"      # final CSV
EXPORT_ALL_PDF_PAGES = False            # True: convert all PDF pages
PDF_DPI = 200                           # PDF render DPI
# =====================

def ensure_outdir(path: Path):
    path.mkdir(parents=True, exist_ok=True)

def to_png_from_image(src_path: Path, out_dir: Path) -> Path:
    with Image.open(src_path) as im:
        if im.mode not in ("RGB", "L"):
            im = im.convert("RGB")
        out_path = out_dir / f"{src_path.stem}.png"
        im.save(out_path, format="PNG")
        return out_path

def to_png_from_pdf(src_path: Path, out_dir: Path, all_pages: bool, dpi: int) -> List[Path]:
    doc = fitz.open(src_path)
    out_paths = []
    zoom = dpi / 72.0
    mat = fitz.Matrix(zoom, zoom)
    pages = range(len(doc)) if all_pages else range(1)
    for i in pages:
        pix = doc.load_page(i).get_pixmap(matrix=mat, alpha=False)
        out_path = out_dir / (f"{src_path.stem}_p{i}.png" if all_pages else f"{src_path.stem}.png")
        pix.save(out_path.as_posix())
        out_paths.append(out_path)
    doc.close()
    return out_paths

def process_row(rel_source: str, image_root: Path, out_dir: Path) -> Tuple[str, str, str]:
    full_path = (image_root / rel_source).resolve()
    png_field = ""
    try:
        if full_path.suffix.lower() in [".png", ".jpg", ".jpeg"]:
            out_path = to_png_from_image(full_path, out_dir)
            png_field = out_path.as_posix()
        elif full_path.suffix.lower() == ".pdf":
            outs = to_png_from_pdf(full_path, out_dir, EXPORT_ALL_PDF_PAGES, PDF_DPI)
            png_field = ";".join(p.as_posix() for p in outs)
        else:
            print(f"[WARN] Unsupported file type: {rel_source}", file=sys.stderr)
    except Exception as e:
        print(f"[ERROR] Failed to process {rel_source}: {e}", file=sys.stderr)
    return (rel_source, png_field)

def main():
    parser = argparse.ArgumentParser(description="Convert relative paths in CSV to PNGs using image root dir")
    parser.add_argument("--csv", required=True, help="CSV file with Source and Caption columns")
    parser.add_argument("--image-root", required=True, help="Root directory where image files are located")
    args = parser.parse_args()

    csv_path = Path(args.csv)
    image_root = Path(args.image_root).resolve()
    out_dir = Path(OUTPUT_DIR)
    out_csv = Path(OUTPUT_CSV)

    if not csv_path.exists():
        print(f"[FATAL] CSV not found: {csv_path}", file=sys.stderr)
        sys.exit(1)
    if not image_root.exists():
        print(f"[FATAL] Image root not found: {image_root}", file=sys.stderr)
        sys.exit(1)

    ensure_outdir(out_dir)
    rows = []

    with csv_path.open("r", encoding="utf-8-sig", newline="") as f:
        reader = csv.DictReader(f)
        if "Source" not in reader.fieldnames:
            print("[FATAL] CSV must contain 'Source' and 'Caption' columns", file=sys.stderr)
            sys.exit(1)
        for row in reader:
            rel_src = row["Source"].strip()

            if not rel_src:
                rows.append((rel_src, ""))
                continue
            rows.append(process_row(rel_src,image_root, out_dir))

    with out_csv.open("w", encoding="utf-8", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["Source", "PNG"])
        writer.writerows(rows)

    print(f"\n✅ PNGs saved to: {out_dir.resolve()}")
    print(f"✅ Output CSV saved to: {out_csv.resolve()}")

if __name__ == "__main__":
    main()
